package com.wt.admin.service.ai.impl;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.ObjectUtil;
import com.wt.admin.domain.dto.ai.AgentsInfoDTO;
import com.wt.admin.domain.dto.ai.ChatDTO;
import com.wt.admin.domain.dto.ai.ChatModelContentDTO;
import com.wt.admin.domain.entity.ai.AgentsInfoEntity;
import com.wt.admin.domain.entity.ai.ChatModelContentEntity;
import com.wt.admin.domain.vo.ai.AgentsInfoVO;
import com.wt.admin.domain.vo.ai.ChatModelContentVO;
import com.wt.admin.domain.vo.sys.UserVO;
import com.wt.admin.service.ai.AgentsInfoService;
import com.wt.admin.service.ai.ChatModelContentService;
import com.wt.admin.service.ai.ChatProxyService;
import com.wt.admin.service.ai.impl.agents.AgentsManager;
import com.wt.admin.service.vector.Vector;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.stream.Collectors;


@Service("chatProxy")
@Slf4j
public class ChatProxyServiceImpl implements ChatProxyService {

    @Autowired
    private Vector vector;
    @Resource
    private ChatModelContentService chatModelContentService;
    @Resource
    ChatMemory chatMemory;
    @Resource
    AgentsInfoService agentsInfoService;
    @Resource
    AgentsManager agentsManager;


    @Override
    public ChatModelContentService getChatModelContentService() {
        return chatModelContentService;
    }

    @Override
    public ChatClient.ChatClientRequestSpec question(ChatDTO dto, UserVO user) {
        AgentsInfoEntity.Config config = BeanUtil.copyProperties(dto.getAgents().getConfig(), AgentsInfoEntity.Config.class);
        AgentsManager.ChatClientMapper chatClientMapper = agentsManager.get(dto.getAgentsId());
        ChatClient.ChatClientRequestSpec chatClient = chatClientMapper.getChatClient().prompt();
        String prompt = dto.getChatConfig().getPrompt();
//        List<Message> messages = chatMemory.get(dto.getContentId().toString());
//        chatMemory.add(dto.getContentId().toString(),messages);
        if (ObjectUtil.isNotEmpty(prompt)) {
            prompt = prompt.replace("{name}", dto.getAgents().getName());
            chatClient = chatClientMapper.getChatClient().prompt(prompt);
        }
        if (CollUtil.isNotEmpty(config.getKnowledges())) {
            String knowledgesStr = config.getKnowledges().stream()
                    .map(Object::toString)
                    .collect(Collectors.joining(","));
            SearchRequest build = SearchRequest.builder().query(dto.getMessage())
                    .filterExpression("knowledgeTitleId in [" + knowledgesStr + "]")
                    .topK(dto.getChatConfig().getTopK())
                    .similarityThreshold(dto.getChatConfig().getSimilarityThreshold())
                    .build();
            List<Document> documents = vector.getElasticsearchVectorStore().similaritySearch(build);
            log.debug("找到相似文档 {}", documents.size());
            chatClient.advisors(QuestionAnswerAdvisor.builder(vector.getElasticsearchVectorStore()).searchRequest(build).build());
        }
        chatClient.advisors(MessageChatMemoryAdvisor.builder(chatMemory).conversationId(dto.getContentId().toString()).build());
        if(dto.getChatConfig().getEnableSynergism()){
            chatClient = synergism(chatClient,config.getSynergism(),dto.getMessage());
        }
        return chatClient.user(dto.getMessage());
    }

    @Override
    public ChatModelContentVO modeContentDel(ChatModelContentDTO data) {
        ChatModelContentVO chatModelContentVO = chatModelContentService.modeContentDel(data);
        chatMemory.clear(data.getId().toString());
        return chatModelContentVO;
    }

    @Override
    public ChatModelContentVO modeContentEdit(ChatDTO data, UserVO user) {
        ChatModelContentVO chatModelContentVO = chatModelContentService.modeContentEdit(data, user);
        // 根据model_id查询
//        chatMemory.add(chatModelContentVO.getId().toString(),new UserMessage(data.getMessage()));
        return chatModelContentVO;
    }

    @Override
    public ChatModelContentVO modeContentInfo(ChatModelContentDTO data, UserVO user) {
        ChatModelContentVO chatModelContentVO = chatModelContentService.modeContentInfo(data, user);
        chatMemory.clear(chatModelContentVO.getId().toString());
        List<Message> list = chatModelContentVO.getContent().stream().map(i -> BeanUtil.copyProperties(i, ChatModelContentEntity.Contents.class))
                .map(ChatModelContentEntity.Contents::toMessage).toList();
        chatMemory.add(chatModelContentVO.getId().toString(), list);
        return chatModelContentVO;
    }

    @Override
    public void reply(String string, ChatDTO data, UserVO user) {
        chatModelContentService.reply(string, data, user);
//        chatMemory.add(data.getContentId().toString(), new SystemMessage(string));
    }

    @Override
    public List<AgentsInfoVO> agentsList(AgentsInfoDTO data) {
        return agentsInfoService.infoList(data);
    }

    private ChatClient.ChatClientRequestSpec synergism(ChatClient.ChatClientRequestSpec chat, List<List<AgentsInfoEntity.AgentSynergism>> synergisms,String msg) {
        if(CollUtil.isEmpty(synergisms)){
            return chat;
        }
        List<String> list = new ArrayList<>();
        for (List<AgentsInfoEntity.AgentSynergism> synergism : synergisms) {
            List<String> context = synergism.stream().map(agent -> CompletableFuture.supplyAsync(() -> {
                ChatClient chatClient = agentsManager.get(agent.agent).getChatClient();
                ChatClient.ChatClientRequestSpec prompt = chatClient.prompt(agent.getPrompt());
                if(!list.isEmpty()){
                    prompt = prompt.system(String.join("\n", list));
                }
                return prompt.user(msg).call().content();
            })).toList().stream().map(CompletableFuture::join).toList();
            list.clear();
            list.addAll(context);
        }
        // 更新 chat 内容
        return chat.system(String.join("\n", list));
    }

}
